#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
股票买卖点预测算法
基于技术分析指标的买卖信号生成系统
"""

import numpy as np
import pandas as pd
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from enum import Enum
import json

class Signal(Enum):
    """交易信号"""
    STRONG_BUY = "强烈买入"
    BUY = "买入"
    HOLD = "持有"
    SELL = "卖出"
    STRONG_SELL = "强烈卖出"

@dataclass
class TradingPoint:
    """交易点"""
    date: str
    price: float
    signal: Signal
    confidence: float  # 信心度 0-100
    reason: str
    indicators: Dict[str, float]

class TechnicalAnalysis:
    """技术分析指标计算"""
    
    @staticmethod
    def calculate_ma(prices: List[float], period: int) -> List[float]:
        """计算移动平均线 MA"""
        if len(prices) < period:
            return []
        
        result = []
        for i in range(len(prices)):
            if i < period - 1:
                result.append(np.nan)
            else:
                ma = np.mean(prices[i-period+1:i+1])
                result.append(ma)
        return result
    
    @staticmethod
    def calculate_ema(prices: List[float], period: int) -> List[float]:
        """计算指数移动平均线 EMA"""
        if len(prices) < period:
            return []
        
        result = []
        multiplier = 2 / (period + 1)
        
        # 第一个EMA使用SMA
        sma = np.mean(prices[:period])
        result.append(sma)
        
        # 后续使用EMA公式
        for i in range(1, len(prices)):
            if i < period:
                result.append(np.nan)
            else:
                ema = (prices[i] - result[-1]) * multiplier + result[-1]
                result.append(ema)
        
        return result
    
    @staticmethod
    def calculate_macd(prices: List[float], fast=12, slow=26, signal=9) -> Dict[str, List[float]]:
        """计算MACD指标"""
        ema_fast = TechnicalAnalysis.calculate_ema(prices, fast)
        ema_slow = TechnicalAnalysis.calculate_ema(prices, slow)
        
        macd_line = [fast - slow if not np.isnan(fast) and not np.isnan(slow) else np.nan 
                     for fast, slow in zip(ema_fast, ema_slow)]
        
        signal_line = TechnicalAnalysis.calculate_ema(
            [x for x in macd_line if not np.isnan(x)], signal
        )
        
        # 补齐长度
        signal_line = [np.nan] * (len(macd_line) - len(signal_line)) + signal_line
        
        histogram = [m - s if not np.isnan(m) and not np.isnan(s) else np.nan 
                     for m, s in zip(macd_line, signal_line)]
        
        return {
            'macd': macd_line,
            'signal': signal_line,
            'histogram': histogram
        }
    
    @staticmethod
    def calculate_rsi(prices: List[float], period: int = 14) -> List[float]:
        """计算相对强弱指标 RSI"""
        if len(prices) < period + 1:
            return []
        
        deltas = [prices[i] - prices[i-1] for i in range(1, len(prices))]
        
        gains = [d if d > 0 else 0 for d in deltas]
        losses = [-d if d < 0 else 0 for d in deltas]
        
        result = [np.nan]  # 第一个价格无法计算
        
        avg_gain = np.mean(gains[:period])
        avg_loss = np.mean(losses[:period])
        
        for i in range(period, len(deltas)):
            if avg_loss == 0:
                rsi = 100
            else:
                rs = avg_gain / avg_loss
                rsi = 100 - (100 / (1 + rs))
            
            result.append(rsi)
            
            # 更新平均值
            avg_gain = (avg_gain * (period - 1) + gains[i]) / period
            avg_loss = (avg_loss * (period - 1) + losses[i]) / period
        
        return result
    
    @staticmethod
    def calculate_bollinger_bands(prices: List[float], period: int = 20, num_std: float = 2) -> Dict[str, List[float]]:
        """计算布林带"""
        ma = TechnicalAnalysis.calculate_ma(prices, period)
        
        upper = []
        lower = []
        
        for i in range(len(prices)):
            if i < period - 1:
                upper.append(np.nan)
                lower.append(np.nan)
            else:
                std = np.std(prices[i-period+1:i+1])
                upper.append(ma[i] + num_std * std)
                lower.append(ma[i] - num_std * std)
        
        return {
            'upper': upper,
            'middle': ma,
            'lower': lower
        }
    
    @staticmethod
    def calculate_kdj(highs: List[float], lows: List[float], closes: List[float], 
                      n: int = 9, m1: int = 3, m2: int = 3) -> Dict[str, List[float]]:
        """计算KDJ指标"""
        k_values = []
        d_values = []
        j_values = []
        
        rsv_values = []
        for i in range(len(closes)):
            if i < n - 1:
                rsv_values.append(np.nan)
            else:
                low_n = min(lows[i-n+1:i+1])
                high_n = max(highs[i-n+1:i+1])
                if high_n == low_n:
                    rsv = 50
                else:
                    rsv = (closes[i] - low_n) / (high_n - low_n) * 100
                rsv_values.append(rsv)
        
        # K值
        k = 50
        d = 50
        for rsv in rsv_values:
            if np.isnan(rsv):
                k_values.append(np.nan)
                d_values.append(np.nan)
                j_values.append(np.nan)
            else:
                k = (2 * k + rsv) / 3
                d = (2 * d + k) / 3
                j = 3 * k - 2 * d
                
                k_values.append(k)
                d_values.append(d)
                j_values.append(j)
        
        return {
            'k': k_values,
            'd': d_values,
            'j': j_values
        }

class BuySellPredictor:
    """买卖点预测器"""
    
    def __init__(self):
        self.ta = TechnicalAnalysis()
    
    def analyze(self, stock_data: List[Dict]) -> List[TradingPoint]:
        """
        分析股票数据，生成买卖点
        
        stock_data格式:
        [
            {'date': '2025-01-01', 'open': 100, 'close': 105, 'high': 106, 'low': 99, 'volume': 1000},
            ...
        ]
        """
        if len(stock_data) < 30:
            return []
        
        # 提取数据
        dates = [d['date'] for d in stock_data]
        opens = [float(d['open']) for d in stock_data]
        closes = [float(d['close']) for d in stock_data]
        highs = [float(d['high']) for d in stock_data]
        lows = [float(d['low']) for d in stock_data]
        volumes = [float(d['volume']) for d in stock_data]
        
        # 计算所有指标
        ma5 = self.ta.calculate_ma(closes, 5)
        ma10 = self.ta.calculate_ma(closes, 10)
        ma20 = self.ta.calculate_ma(closes, 20)
        ma60 = self.ta.calculate_ma(closes, 60) if len(closes) >= 60 else [np.nan] * len(closes)
        
        macd = self.ta.calculate_macd(closes)
        rsi = self.ta.calculate_rsi(closes, 14)
        bollinger = self.ta.calculate_bollinger_bands(closes, 20)
        kdj = self.ta.calculate_kdj(highs, lows, closes)
        
        # 生成交易信号
        trading_points = []
        
        for i in range(20, len(stock_data)):  # 从第20天开始分析，确保有足够数据
            signals = []
            reasons = []
            
            # 1. MACD金叉/死叉
            if i > 0 and not np.isnan(macd['macd'][i]) and not np.isnan(macd['signal'][i]):
                if macd['macd'][i-1] <= macd['signal'][i-1] and macd['macd'][i] > macd['signal'][i]:
                    signals.append(2)  # 买入信号
                    reasons.append("MACD金叉")
                elif macd['macd'][i-1] >= macd['signal'][i-1] and macd['macd'][i] < macd['signal'][i]:
                    signals.append(-2)  # 卖出信号
                    reasons.append("MACD死叉")
            
            # 2. RSI超买超卖
            if i < len(rsi) and not np.isnan(rsi[i]):
                if rsi[i] < 30:
                    signals.append(2)
                    reasons.append(f"RSI超卖({rsi[i]:.1f})")
                elif rsi[i] > 70:
                    signals.append(-2)
                    reasons.append(f"RSI超买({rsi[i]:.1f})")
            
            # 3. 布林带突破
            if not np.isnan(bollinger['lower'][i]) and not np.isnan(bollinger['upper'][i]):
                if closes[i] < bollinger['lower'][i]:
                    signals.append(1)
                    reasons.append("触及下轨")
                elif closes[i] > bollinger['upper'][i]:
                    signals.append(-1)
                    reasons.append("触及上轨")
            
            # 4. 均线系统
            if not np.isnan(ma5[i]) and not np.isnan(ma10[i]) and not np.isnan(ma20[i]):
                # 金叉：短期均线上穿长期均线
                if i > 0:
                    if ma5[i-1] <= ma10[i-1] and ma5[i] > ma10[i]:
                        signals.append(2)
                        reasons.append("MA5上穿MA10")
                    elif ma5[i-1] >= ma10[i-1] and ma5[i] < ma10[i]:
                        signals.append(-2)
                        reasons.append("MA5下穿MA10")
                
                # 多头排列
                if ma5[i] > ma10[i] > ma20[i]:
                    signals.append(1)
                    reasons.append("多头排列")
                # 空头排列
                elif ma5[i] < ma10[i] < ma20[i]:
                    signals.append(-1)
                    reasons.append("空头排列")
            
            # 5. KDJ指标
            if not np.isnan(kdj['k'][i]) and not np.isnan(kdj['d'][i]):
                if kdj['k'][i] < 20 and kdj['d'][i] < 20:
                    signals.append(1)
                    reasons.append(f"KDJ超卖({kdj['k'][i]:.1f})")
                elif kdj['k'][i] > 80 and kdj['d'][i] > 80:
                    signals.append(-1)
                    reasons.append(f"KDJ超买({kdj['k'][i]:.1f})")
                
                # KDJ金叉死叉
                if i > 0 and not np.isnan(kdj['k'][i-1]):
                    if kdj['k'][i-1] <= kdj['d'][i-1] and kdj['k'][i] > kdj['d'][i] and kdj['k'][i] < 50:
                        signals.append(2)
                        reasons.append("KDJ金叉")
                    elif kdj['k'][i-1] >= kdj['d'][i-1] and kdj['k'][i] < kdj['d'][i] and kdj['k'][i] > 50:
                        signals.append(-2)
                        reasons.append("KDJ死叉")
            
            # 6. 成交量分析
            if i >= 5:
                avg_volume = np.mean(volumes[i-5:i])
                if volumes[i] > avg_volume * 1.5:
                    if closes[i] > opens[i]:
                        signals.append(1)
                        reasons.append("放量上涨")
                    else:
                        signals.append(-1)
                        reasons.append("放量下跌")
            
            # 综合信号强度
            if signals:
                signal_sum = sum(signals)
                confidence = min(abs(signal_sum) * 15, 95)  # 信心度
                
                # 判断最终信号
                if signal_sum >= 4:
                    final_signal = Signal.STRONG_BUY
                elif signal_sum >= 2:
                    final_signal = Signal.BUY
                elif signal_sum <= -4:
                    final_signal = Signal.STRONG_SELL
                elif signal_sum <= -2:
                    final_signal = Signal.SELL
                else:
                    final_signal = Signal.HOLD
                
                # 只记录买入和卖出信号
                if final_signal != Signal.HOLD:
                    trading_point = TradingPoint(
                        date=dates[i],
                        price=closes[i],
                        signal=final_signal,
                        confidence=confidence,
                        reason=" | ".join(reasons[:3]),  # 最多显示3个原因
                        indicators={
                            'ma5': ma5[i] if i < len(ma5) and not np.isnan(ma5[i]) else 0,
                            'ma10': ma10[i] if i < len(ma10) and not np.isnan(ma10[i]) else 0,
                            'ma20': ma20[i] if i < len(ma20) and not np.isnan(ma20[i]) else 0,
                            'rsi': rsi[i] if i < len(rsi) and not np.isnan(rsi[i]) else 0,
                            'macd': macd['macd'][i] if i < len(macd['macd']) and not np.isnan(macd['macd'][i]) else 0,
                            'kdj_k': kdj['k'][i] if i < len(kdj['k']) and not np.isnan(kdj['k'][i]) else 0,
                        }
                    )
                    trading_points.append(trading_point)
        
        return trading_points
    
    def get_current_recommendation(self, stock_data: List[Dict]) -> Dict:
        """获取当前的交易建议"""
        if len(stock_data) < 30:
            return {
                'signal': 'HOLD',
                'confidence': 0,
                'reason': '数据不足',
                'price': 0
            }
        
        trading_points = self.analyze(stock_data)
        
        if not trading_points:
            return {
                'signal': 'HOLD',
                'confidence': 50,
                'reason': '无明显信号',
                'price': stock_data[-1]['close']
            }
        
        # 返回最新的信号
        latest = trading_points[-1]
        return {
            'signal': latest.signal.value,
            'confidence': latest.confidence,
            'reason': latest.reason,
            'price': latest.price,
            'date': latest.date,
            'indicators': latest.indicators
        }
    
    def backtest(self, stock_data: List[Dict], initial_capital: float = 100000) -> Dict:
        """回测交易策略"""
        trading_points = self.analyze(stock_data)
        
        capital = initial_capital
        shares = 0
        trades = []
        
        for point in trading_points:
            if point.signal in [Signal.STRONG_BUY, Signal.BUY] and capital > 0:
                # 买入
                shares_to_buy = int(capital / point.price)
                if shares_to_buy > 0:
                    cost = shares_to_buy * point.price
                    capital -= cost
                    shares += shares_to_buy
                    trades.append({
                        'date': point.date,
                        'action': 'BUY',
                        'price': point.price,
                        'shares': shares_to_buy,
                        'reason': point.reason
                    })
            
            elif point.signal in [Signal.STRONG_SELL, Signal.SELL] and shares > 0:
                # 卖出
                revenue = shares * point.price
                capital += revenue
                trades.append({
                    'date': point.date,
                    'action': 'SELL',
                    'price': point.price,
                    'shares': shares,
                    'reason': point.reason
                })
                shares = 0
        
        # 计算最终价值
        final_price = stock_data[-1]['close']
        final_value = capital + shares * final_price
        profit = final_value - initial_capital
        profit_rate = (profit / initial_capital) * 100
        
        # 计算买入持有策略收益
        buy_hold_shares = int(initial_capital / stock_data[0]['close'])
        buy_hold_value = buy_hold_shares * final_price
        buy_hold_profit = buy_hold_value - initial_capital
        buy_hold_rate = (buy_hold_profit / initial_capital) * 100
        
        return {
            'initial_capital': initial_capital,
            'final_value': final_value,
            'profit': profit,
            'profit_rate': profit_rate,
            'trades_count': len(trades),
            'trades': trades,
            'buy_hold_rate': buy_hold_rate,
            'excess_return': profit_rate - buy_hold_rate
        }

def main():
    """测试预测算法"""
    # 模拟股票数据
    print("="*80)
    print("📊 股票买卖点预测算法测试")
    print("="*80)
    
    # 生成模拟数据（也可以从API获取真实数据）
    np.random.seed(42)
    dates = pd.date_range(start='2024-01-01', periods=100, freq='D')
    
    # 模拟价格走势
    base_price = 100
    prices = [base_price]
    for i in range(99):
        change = np.random.randn() * 2
        new_price = prices[-1] + change
        prices.append(max(new_price, 1))  # 确保价格为正
    
    stock_data = []
    for i, date in enumerate(dates):
        stock_data.append({
            'date': date.strftime('%Y-%m-%d'),
            'open': prices[i] * (1 + np.random.uniform(-0.02, 0.02)),
            'close': prices[i],
            'high': prices[i] * (1 + abs(np.random.uniform(0, 0.03))),
            'low': prices[i] * (1 - abs(np.random.uniform(0, 0.03))),
            'volume': np.random.randint(10000, 100000)
        })
    
    # 创建预测器
    predictor = BuySellPredictor()
    
    # 分析买卖点
    print("\n📈 分析买卖点...")
    trading_points = predictor.analyze(stock_data)
    
    print(f"\n找到 {len(trading_points)} 个交易信号:\n")
    print(f"{'日期':<12} {'价格':<10} {'信号':<12} {'信心度':<8} {'原因':<40}")
    print("-" * 85)
    
    for point in trading_points[-10:]:  # 显示最后10个信号
        print(f"{point.date:<12} {point.price:<10.2f} {point.signal.value:<12} "
              f"{point.confidence:<8.1f} {point.reason:<40}")
    
    # 当前建议
    print("\n" + "="*80)
    print("💡 当前交易建议")
    print("="*80)
    
    recommendation = predictor.get_current_recommendation(stock_data)
    print(f"\n信号: {recommendation['signal']}")
    print(f"信心度: {recommendation['confidence']:.1f}%")
    print(f"价格: ¥{recommendation['price']:.2f}")
    print(f"原因: {recommendation['reason']}")
    
    if 'indicators' in recommendation:
        print(f"\n技术指标:")
        for key, value in recommendation['indicators'].items():
            print(f"  {key.upper()}: {value:.2f}")
    
    # 回测
    print("\n" + "="*80)
    print("🔬 策略回测")
    print("="*80)
    
    backtest_result = predictor.backtest(stock_data, 100000)
    
    print(f"\n初始资金: ¥{backtest_result['initial_capital']:,.2f}")
    print(f"最终价值: ¥{backtest_result['final_value']:,.2f}")
    print(f"盈亏金额: ¥{backtest_result['profit']:,.2f}")
    print(f"收益率: {backtest_result['profit_rate']:.2f}%")
    print(f"交易次数: {backtest_result['trades_count']}")
    print(f"\n对比买入持有策略:")
    print(f"买入持有收益率: {backtest_result['buy_hold_rate']:.2f}%")
    print(f"超额收益: {backtest_result['excess_return']:.2f}%")
    
    print(f"\n最近5笔交易:")
    print(f"{'日期':<12} {'操作':<6} {'价格':<10} {'数量':<10} {'原因':<40}")
    print("-" * 85)
    for trade in backtest_result['trades'][-5:]:
        print(f"{trade['date']:<12} {trade['action']:<6} {trade['price']:<10.2f} "
              f"{trade['shares']:<10} {trade['reason']:<40}")
    
    # 导出JSON格式
    print("\n" + "="*80)
    print("💾 导出预测结果")
    print("="*80)
    
    output = {
        'trading_points': [
            {
                'date': p.date,
                'price': p.price,
                'signal': p.signal.value,
                'confidence': p.confidence,
                'reason': p.reason,
                'indicators': p.indicators
            }
            for p in trading_points
        ],
        'current_recommendation': recommendation,
        'backtest': backtest_result
    }
    
    with open('prediction_result.json', 'w', encoding='utf-8') as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    
    print("\n✅ 结果已保存到 prediction_result.json")
    print("\n" + "="*80)

if __name__ == "__main__":
    main()

